################################################################################
# Plot.py - This module contains a function for plotting lists of spectra
#
"""
TODO: Fix memory leak for list of files?
"""

__author__ = 'Adrian Price-Whelan <adrn@astro.columbia.edu>'

# Standard library dependencies
import sys, os
import logging

# Third party dependencies
import numpy as np
import pyfits as pf
import matplotlib as ml
ml.rcParams['xtick.labelsize'] = 8
ml.rcParams['ytick.labelsize'] = 8
import matplotlib.pyplot as plt
    
try:
    import astropysics.spec as astropysics_spec
    astropysics_spectrum = astropysics_spec.Spectrum([0,1], [10,20])
    astropysics_loaded = True
except ImportError:
    # warnings.warn()
    astropysics_loaded = False


def plotspec(spectra, use_inverse_variance=False, shift_to_rest_frame=False, lines=[], **kwargs):
    """ Accepts a list of generic Spectrum objects, a list of FITS file names, 
        or a list of astropysics spectrum objects and plots the spectra using matplotlib. 
        Note, it also accepts a **mixture** of these types! See below for more information 
        on each of these formats.
        
        You may also pass a list of Line() objects (see Lines.py), or a LineGroup()
        object in to this function using the keyword 'lines'. These lines get plotted
        along with the spectrum.
        
        The default behavior of this function is to plot the spectra on their own figures, and
        save them as numbered png's. There are a lot of other optional possibilities, so see
        below in "Optional kwargs."
        
        Input Parameters:
        --------------------
        spectra : 
            Spectrum objects:
                + 'spectra' can be a list of any custom object that represents a spectrum, 
                such that it has a few required attributes:
                    - flux
                    - wavelengths
                    - inverse_variance (optional, unless you set use_inverse_variance == True)
                    - z (redshift)
            
            FITS file names:
                + Right now, the only format of FITS file accepted by this function is to
                have a TableHDU as the 1st HDU (not the Primary HDU / 0th) with columns
                'flux', 'wavelengths', and optionally 'inverse_variance' or 'ivar'
            
            Tuple:
                + Must be either (wavelengths, flux) or (wavelengths, flux, inverse variance)
                + With this format, you can't shift to the rest frame (no Z supplied), but I'll
                  eventually add that as a kwarg
            
            astropysics.spec.Spectrum:
                + Coming Soon
        
        use_inverse_variance : boolean (default=False)
            Specify whether to look for an inverse_variance array, and weight the spectrum
            using these errors
        
        lines : list of Line() objects (see Lines.py)
            If lines are specified, they will be plotted as vertical lines along with the spectra.
        
        shift_to_rest_frame : boolean (default=True)
            Specify whether to look for a redshift, and shift the spectrum to the rest frame
        
        Optional kwargs:
        ----------------
            axes : matplotlib Axes object, or a list of axes objects, or a matplotlib figure object
                If a single matplotlib Axes object is passed, it will plot all spectra on the same 
                axis. If a list of axes objects is passed, it must have the same length as the list
                of spectra, and the list (with one spectrum plotted per axis) is returned. If a 
                figure object is passed, it will plot one spectrum per row -- careful with this one
                if you have > 3 spectra, because the plots will be HUGE.
            
            save : boolean
                If false, it will show() each spectrum. Not recommended for large numbers of
                spectra!
            
            show : boolean
                If both save and show are False, the function will just return the axes object.
            
            filename : string (e.g. "plot-1234"), or list of strings
                With or without the extension (it will strip off the extension if one is found,
                and use that for the filetype). If a list, must be of length len(spectra)!
                
            filetype : string (e.g. "pdf"), or list of strings
                If a list, must be of length len(spectra)!
            
            title : string, or list of strings
                If a list, must be of length len(spectra)!
            
            xlabel : string, or list of strings
                If a list, must be of length len(spectra)!
            
            ylabel : string, or list of strings
                If a list, must be of length len(spectra)!
            
            xlim : tuple, or list of tuples
                If a list, must be of length len(spectra)!
            
            ylim : tuple, or list of tuples
                If a list, must be of length len(spectra)!
            
    """
    
    if not isinstance(spectra, [].__class__):
        spectra = [spectra]
    
    if not isinstance(lines, [].__class__):
        lines = [lines]
    
    if kwargs.has_key("axes"):
        if isinstance(kwargs["axes"], [].__class__) and len(kwargs["axes"]) != len(spectra):
            raise ValueError("Length of 'axes' must be same as 'spectra'!")
        elif isinstance(kwargs["axes"], [].__class__):
            axes = kwargs["axes"]
            createAxes = False
        else:
            axes = [kwargs["axes"]] * len(spectra)
            createAxes = False
        del kwargs['axes']
    else:
        axes = []
        createAxes = True
    
    if kwargs.has_key("filename"):
        if isinstance(kwargs["filename"], [].__class__) and len(kwargs["filename"]) != len(spectra):
            raise ValueError("Length of 'filename' must be same as 'spectra'!")
        elif isinstance(kwargs["filename"], [].__class__):
            filenames = kwargs["filename"]
        else:
            filenames = [kwargs["filename"]]
            if len(filenames) != len(spectra):
                raise ValueError("Length of 'filename' must be same as 'spectra'! You must specify more than one filename.")
    else:
        filenames = None
    
    if kwargs.has_key("filetype"):
        if isinstance(kwargs["filetype"], [].__class__) and len(kwargs["filetype"]) != len(spectra):
            raise ValueError("Length of 'filetype' must be same as 'spectra'!")
        elif isinstance(kwargs["filetype"], [].__class__):
            filetypes = kwargs["filetype"]
        else:
            filetypes = [kwargs["filetype"]]*len(spectra)
    else:
        filetypes = None
    
    if kwargs.has_key("title"):
        if isinstance(kwargs["title"], [].__class__) and len(kwargs["title"]) != len(spectra):
            raise ValueError("Length of 'title' must be same as 'spectra'!")
        elif isinstance(kwargs["title"], [].__class__):
            titles = kwargs["title"]
        else:
            titles = [kwargs["title"]]*len(spectra)
    else:
        titles = None
        
    if kwargs.has_key("xlabel"):
        if isinstance(kwargs["xlabel"], [].__class__) and len(kwargs["xlabel"]) != len(spectra):
            raise ValueError("Length of 'xlabel' must be same as 'spectra'!")
        elif isinstance(kwargs["xlabel"], [].__class__):
            xlabels = kwargs["xlabel"]
        else:
            xlabels = [kwargs["xlabel"]]*len(spectra)
    else:
        xlabels = None
    
    if kwargs.has_key("ylabel"):
        if isinstance(kwargs["ylabel"], [].__class__) and len(kwargs["ylabel"]) != len(spectra):
            raise ValueError("Length of 'ylabel' must be same as 'spectra'!")
        elif isinstance(kwargs["ylabel"], [].__class__):
            ylabels = kwargs["ylabel"]
        else:
            ylabels = [kwargs["ylabel"]]*len(spectra)
    else:
        ylabels = None
    
    if kwargs.has_key('save'):
        savePlots = kwargs['save']
        del kwargs['save']
    else:
        savePlots = False
    
    if kwargs.has_key('show'):
        showPlots = kwargs['show']
        del kwargs['show']
    else:
        showPlots = False
    
    for ii,spectrum in enumerate(spectra):
        logging.debug("Spectrum %d" % ii)
        
        # Determine whether spectrum is a generic object, filename, or astropysics object
        if isinstance(spectrum, "".__class__):
            logging.debug("Spectrum is a filename at %s" % spectrum)
            
            # filename
            if os.path.exists(spectrum):
                try:
                    hduList = pf.open(spectrum)
                except IOError:
                    raise IOError("Invalid FITS file! %s" % spectrum)
            else:
                raise IOError("File not found! %s" % spectrum)
            
            flux = hduList[1].data.field("flux")
            wavelengths = hduList[1].data.field("wavelengths")
            
            # If the user wants to weight the flux using the inverse variance, make
            #   sure the inv. var. is in the file!
            if use_inverse_variance:
                if "ivar" in hduList[1].data.names:
                    inverse_variance = hduList[1].data.field("ivar")
                elif "inverse_variance" in hduList[1].data.names:
                    inverse_variance = hduList[1].data.field("inverse_variance")
                else:
                    raise ValueError("You specified 'use_inverse_variance', but the inverse variance could not be found in %s!" % spectrum)
                logging.debug("Inverse variance found for spectrum %d" % ii)
            
            try:
                z = hduList[0].header['Z']
            except KeyError:
                try:
                    z = hduList[1].header['Z']
                except KeyError:
                    raise ValueError("The redshift could not be found in %s!" % spectrum)
                    
        elif astropysics_loaded and isinstance(spectrum, astropysics_spectrum.__class__):
            # astropysics Spectrum object
            print "Currently Not Supported! Coming soon..."
            sys.exit(1)
        
        elif isinstance(spectrum, ().__class__):
            logging.debug("Spectrum is a tuple of length %s" % len(spectrum))
            shift_to_rest_frame = False
            
            if len(spectrum) == 2:
                wavelengths, flux = spectrum
            elif len(spectrum) == 3:
                wavelengths, flux, inverse_variance = spectrum
            else:
                raise ValueError("You specified a spectrum as a tuple, but it has an invalid length %d" % len(spectrum))
        
        else:
            logging.debug("Spectrum is some other object of type %s" % spectrum.__class__)
            
            # some other object, check to see if it has flux, wavelengths[, inverse_variance]
            flux = spectrum.flux
            wavelengths = spectrum.wavelengths
            z = spectrum.z
            
            # If the user wants to weight the flux using the inverse variance, make
            #   sure the inv. var. is in the file!
            if use_inverse_variance:
                try:
                    inverse_variance = spectrum.inverse_variance
                except AttributeError:
                    raise AttributeError("You specified 'use_inverse_variance', but the spectrum object has no attribute inverse_variance!")
                logging.debug("Inverse variance found for spectrum %d" % ii)
        
        try:
            logging.debug("Redshift %s" % z)
        except UnboundLocalError:
            logging.debug("Redshift not defined")
        
        # Shift the spectrum to the rest frame
        if shift_to_rest_frame:
            logging.debug("Shifting to rest frame")
            flux = flux * (1. + z)
            wavelengths = wavelengths / (1. + z)
            if use_inverse_variance:
                inverse_variance = inverse_variance * (1. + z)**2
        
        if createAxes:
            fig = plt.figure(figsize=(18,10))
            axis = fig.add_subplot(111)
            axes.append(axis)
        else:
            logging.debug("Axis: %s" % axes[ii])
            axis = axes[ii]
            savePlots = False
        
        axis.plot(wavelengths, flux, **kwargs)
        axis.set_xlim(min(wavelengths), max(wavelengths))
        
        if len(lines) != 0:
            for line in lines:
                if shift_to_rest_frame:
                    line.plot(axis, linestyle="--", alpha=0.5)
                else:
                    line.plot(axis, z=z, linestyle="--", alpha=0.5)
        
        if filenames != None:
            name = filenames[ii]
        else:
            name = "%d" % ii
        
        if filetypes != None:
            saveType = filetypes[ii]
        else:
            saveType = "png"
        
        if titles != None:
            logging.debug("title: %s" % titles[ii])
            axis.set_title(titles[ii])
        
        if xlabels != None:
            logging.debug("xlabel: %s" % xlabels[ii])
            axis.set_xlabel(xlabels[ii])
        
        if ylabels != None:
            logging.debug("ylabel: %s" % ylabels[ii])
            axis.set_ylabel(ylabels[ii])

        if savePlots:
            fig.savefig(name + "." + saveType, format=saveType)
        elif showPlots:
            plt.show()

    return axes

if __name__ == "__main__":
    import Lines
    
    logging.basicConfig(level=logging.DEBUG, format='Debug: %(message)s')
    
    fig = plt.figure(figsize=(18,10))
    ax1 = fig.add_subplot(211)
    ax2 = fig.add_subplot(212)
    axes = [ax1, ax2]
    
    spectra = ["/Users/adrian/Downloads/spec-3841-55572-0865.fits", "/Users/adrian/Downloads/spec-4055-55359-0408.fits"]
    filenames = ["spec-3841-55572-0865", "spec-4055-55359-0408"]
    titles = "Derp"
    filetypes = "pdf"
    xlabels = r"Angstroms ($\AA$)"
    ylabels = r"$\mathrm{Flux} \, [10^{-17} \, \mathrm{erg} \, \mathrm{cm}^{-2} \, \mathrm{A}^{-1}]$"
    plotspec(spectra, use_inverse_variance=True, filename=filenames, filetype=filetypes, axes=axes, title=titles, xlabel=xlabels, ylabel=ylabels)
    plt.show()
    del fig
    
    plotspec(spectra, use_inverse_variance=True, lines=[Lines.BalmerLines])
else:
    logging.basicConfig(level=logging.ERROR, format='apwlib error: %(message)s')

#del np, os, sys, ml, pf, logging, astropysics_spectrum, astropysics_spec, astropysics_loaded